#include "Socket.h"
using namespace std;

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Socket
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      Socket::Socket
--
--	RETURNS:			
--
--	NOTES:
--		Constructor for Socket wrapper class
---------------------------------------------------------------------------------------*/
Socket::Socket()
{

}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Socket
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      Socket::Socket(SOCKET s)
--						SOCKET s - SOCKET to instantiate class with
--
--	RETURNS:			
--
--	NOTES:
--		Construct a socket with a SOCKET parameter
---------------------------------------------------------------------------------------*/
Socket::Socket(SOCKET s)
{
	sock = s;
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Socket
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      Socket::Socket(const Socket& s)
--
--	RETURNS:			
--
--	NOTES:
--		Copy constructor
---------------------------------------------------------------------------------------*/
Socket::Socket(const Socket& s)
{
	sock = s.sock;
	sin = s.sin;
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::operator=
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      Socket& operator=(const Socket& s)
--
--	RETURNS:		Socket& - pointer to this instance
--
--	NOTES:
--		overloaded assignment operator
---------------------------------------------------------------------------------------*/
Socket& Socket::operator=(const Socket& s)
{
	sock = s.sock;
	sin = s.sin;
	return *this;
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::operator==
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      bool operator==(const Socket& s)
--
--	RETURNS:		bool
--
--	NOTES:
--		overloaded comparison operator
---------------------------------------------------------------------------------------*/
bool Socket::operator==(const Socket& s)
{
	if (s.sock == sock)
		return true;
	return false;
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Create
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void Socket::Create(string protocol)
--						string protocol - protocol of the socket either "TCP" or "UDP"
--
--	RETURNS:		void
--
--	NOTES:
--		Creates a socket of protocol type and also sets reuseaddr so the port does not
--		get stuck on disconnect for a long period of time
---------------------------------------------------------------------------------------*/
void Socket::Create(string protocol)
{
	int inet;
	int streamType;

	if (protocol == "TCP")
	{
		inet = AF_INET;
		streamType = SOCK_STREAM;
	}
	else if (protocol == "UDP")
	{
		inet = PF_INET;
		streamType = SOCK_DGRAM;
	}
	else
	{
		printf("invalid protocol");
		exit(1);
	}

	if ((sock = WSASocket(inet, streamType, 0, NULL, 0, WSA_FLAG_OVERLAPPED)) == INVALID_SOCKET)
	{
		char str[80];
		sprintf(str, "sock error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}

	int optval = 1;
	if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (char*)&optval, sizeof optval) != 0)
	{
		char str[80];
		sprintf(str, "setsockopt error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Bind
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void Socket::Bind(int port)
--						int port - port to bind on
--
--	RETURNS:		void
--
--	NOTES:
--		Binds a socket to the specified port
---------------------------------------------------------------------------------------*/
void Socket::Bind(int port)
{
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = INADDR_ANY;
	sin.sin_port = htons(port);

	if (bind(sock, (struct sockaddr *)&sin, sizeof(sin)) < 0)
	{
		char str[80];
		sprintf(str, "bind error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::setMulticastServer
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void setMulticastServer(string ip, int port)
--						string ip - multicast address
--						int port - port to run server on
--
--	RETURNS:		void
--
--	NOTES:
--		Sets this socket up as a multicast server
---------------------------------------------------------------------------------------*/
void Socket::setMulticastServer(string ip, int port)
{
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = inet_addr(ip.c_str());
	sin.sin_port = htons(port);

	int optval = 1;
	if (setsockopt(sock, IPPROTO_IP, IP_MULTICAST_TTL, (char*)&optval, sizeof(optval)) != 0)
	{
		char str[80];
		sprintf(str, "setsockopt error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::setMulticastClient
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void setMulticastClient(string ip, int port)
--						string ip - multicast address
--						int port - port to run socket on
--
--	RETURNS:		void
--
--	NOTES:
--		Sets this socket up to receive multicast data
---------------------------------------------------------------------------------------*/
void Socket::setMulticastClient(string ip)
{
	mc_req.imr_multiaddr.s_addr = inet_addr(ip.c_str());
	mc_req.imr_interface.s_addr = INADDR_ANY;
	
	int optval = 1;
	if (setsockopt(sock, IPPROTO_IP, IP_ADD_MEMBERSHIP, (char*)&mc_req, sizeof(mc_req)) != 0)
	{
		char str[80];
		sprintf(str, "setsockopt error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Listen
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void Listen()
--
--	RETURNS:		void
--
--	NOTES:
--		Instructs the socket to start listening
---------------------------------------------------------------------------------------*/
void Socket::Listen()
{
	if (listen(sock, 5) < 0)
	{
		char str[80];
		sprintf(str, "listen error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Accept
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      Socket Socket::Accept()
--
--	RETURNS:		Socket - new connection we have accepted
--
--	NOTES:
--		Accepts a new connection on the listening socket and returns the new socket
---------------------------------------------------------------------------------------*/
Socket Socket::Accept()
{
	int sin_len = sizeof(sin);
	SOCKET newsock;
	if ((newsock = WSAAccept(sock, (struct sockaddr *)&sin, &sin_len, NULL, NULL)) == -1)
	{
		char str[80];
		sprintf(str, "accept error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
	Socket s(newsock);
	return s;
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Connect
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void Connect(string ip, int port)
--						string ip - ip address to connect to
--						int port - port to connect on
--
--	RETURNS:		void
--
--	NOTES:
--		Connects the socket to the specified ip address and port
---------------------------------------------------------------------------------------*/
void Socket::Connect(string ip, int port)
{
	sin.sin_family = AF_INET;
	sin.sin_port = htons(port);
	sin.sin_addr.s_addr = inet_addr(ip.c_str());
	if (connect(sock, (struct sockaddr *)&sin, sizeof(sin)) == -1)
	{
		char str[80];
		sprintf(str, "connect error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}



/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Recv
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      int Recv(string& buffer, unsigned int length, OVERLAPPED* overlapped)
--						string& buffer - buffer to store bytes into
--						unsigned int length - max length to receive
--						OVERLAPPED* overlapped - pointer to the overlapped structure for async IO
--
--	RETURNS:		int - bytes received
--
--	NOTES:
--		Receives data from a socket using overlapped IO. note that if you are sending binary
--		data you should not use this function as it uses strings. use the overloaded char*
--		version instead.
---------------------------------------------------------------------------------------*/
int Socket::Recv(string& buffer, unsigned int length, OVERLAPPED* overlapped)
{
        int bytesReceived = 0;
	DWORD flags = 0;
	WSABUF buf;
	char* c_strBuffer;
	c_strBuffer = (char*)malloc(length);
	buf.buf = c_strBuffer;
	buf.len = length;
	if ((overlapped->hEvent = WSACreateEvent()) == WSA_INVALID_EVENT)
	{
		char str[80];
		sprintf(str, "sock error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
        if (WSARecv(sock, &buf, 1, (LPDWORD)&bytesReceived, &flags, overlapped, NULL) != 0)
	{
		char str[80];
		if (WSAGetLastError() != WSA_IO_PENDING)
		{
			sprintf(str, "recv error: %d", WSAGetLastError());
			printf("%s", str);
			exit(1);
		}
	}
	WSAWaitForMultipleEvents(1, &overlapped->hEvent, TRUE, INFINITE, FALSE);
	WSAResetEvent(overlapped->hEvent);
	WSAGetOverlappedResult(this->getSocket(), overlapped, (LPDWORD)&bytesReceived, TRUE, (LPDWORD)&flags);
	string s(buf.buf);
	buffer = s;
	return bytesReceived;
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Recv
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      int Recv(char* buffer, unsigned int length, OVERLAPPED* overlapped)
--						char* buffer - buffer to store bytes into
--						unsigned int length - max length to receive
--						OVERLAPPED* overlapped - pointer to the overlapped structure for async IO
--
--	RETURNS:		int - bytes received
--
--	NOTES:
--		Receives data from a socket using overlapped IO. overloaded version that uses char*
--		instead of a string. Good for binary data
---------------------------------------------------------------------------------------*/
int Socket::Recv(char* buffer, unsigned int length, OVERLAPPED* overlapped)
{
	int bytesReceived;
	DWORD flags = 0;
	WSABUF buf;
	buf.buf = buffer;
	buf.len = length;
	if (WSARecv(sock, &buf, 1, (LPDWORD)&bytesReceived, &flags, overlapped, NULL) == SOCKET_ERROR)
	{
		char str[80];
		if (WSAGetLastError() != WSA_IO_PENDING)
		{
			sprintf(str, "recv error: %d", WSAGetLastError());
			printf("%s", str);
			exit(1);
		}
	}

	WSAWaitForMultipleEvents(1, &overlapped->hEvent, TRUE, INFINITE, FALSE);
	WSAResetEvent(overlapped->hEvent);
	WSAGetOverlappedResult(this->getSocket(), overlapped, (LPDWORD)&bytesReceived, FALSE, (LPDWORD)&flags);
	return bytesReceived;
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Recv
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      int Recv(char* buffer, unsigned int length, OVERLAPPED* overlapped)
--						char* buffer - buffer to store bytes into
--						unsigned int length - max length to receive
--						OVERLAPPED* overlapped - pointer to the overlapped structure for async IO
--
--	RETURNS:		int - bytes received
--
--	NOTES:
--		Receives data from a socket using overlapped IO. overloaded version that uses char*
--		instead of a string. Good for binary data
---------------------------------------------------------------------------------------*/
int Socket::RecvFrom(char* buffer, unsigned int length)
{
	int sin_len = sizeof(sin);
	int bytesReceived;
	if (bytesReceived = (recvfrom(sock, buffer, length, 0, (struct sockaddr*)&sin, &sin_len) == SOCKET_ERROR))
	{
		char str[80];
		sprintf(str, "recvfrom error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
	return bytesReceived;
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Send
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void Send(string buffer, unsigned int length)
--						string buffer - the buffer to send
--						length of the buffer
--
--	RETURNS:		void
--
--	NOTES:
--		Sends data specified in buffer to the server. Note that this is not good for
--		sending binary data as it uses a string. use char* overloaded function for binary
--		data.
---------------------------------------------------------------------------------------*/
void Socket::Send(string buffer, unsigned int length)
{
	if (send(sock, buffer.c_str(), length+1, 0) < 0)
	{
		char str[80];
		sprintf(str, "send error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::Send
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void Send(char* buffer, unsigned int length)
--						char* buffer - the buffer to send
--						unsigned int length - length of the buffer
--
--	RETURNS:		void
--
--	NOTES:
--		Sends data specified in buffer to the server. char* overloaded version. Good for
--		sending binary data.
---------------------------------------------------------------------------------------*/
void Socket::Send(char* buffer, unsigned int length)
{
	if (send(sock, buffer, length, 0) < 0)
	{
		char str[80];
		sprintf(str, "send error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::SendTo
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      void SendTo(char* buffer, unsigned int length)
--						char* buffer - the buffer to send
--						unsigned int length - length of the buffer
--
--	RETURNS:		void
--
--	NOTES:
--		Sends data specified in buffer to the server specified in sin (UDP)
---------------------------------------------------------------------------------------*/
void Socket::SendTo(char* buffer, unsigned int length)
{
	int sin_len = sizeof(sin);
	if (sendto(sock, buffer, length, 0, (struct sockaddr*)&sin, sin_len) < 0)
	{
		char str[80];
		sprintf(str, "sendTo error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}
}

/*---------------------------------------------------------------------------------------
--	FUNCTION:	Socket::getSocket
--
--	DATE:			March 22, 2010
--
--	REVISIONS		(Date and Description):
--
--	DESIGNER:		Brad Paugh
--
--	PROGRAMMER:		Brad Paugh
--
--  INTERFACE:      SOCKET getSocket()
--
--	RETURNS:		SOCKET - returns the SOCKET used in the instance of this class
--
--	NOTES:
--		returns SOCKET data member
---------------------------------------------------------------------------------------*/
SOCKET Socket::getSocket()
{
	return sock;
}

Socket::~Socket()
{
	/*if (setsockopt(sock, IPPROTO_IP, IP_DROP_MEMBERSHIP, (char*)&mc_req, sizeof(mc_req)) != 0)
	{
		char str[80];
		sprintf(str, "setsockopt error: %d", WSAGetLastError());
		printf("%s", str);
		exit(1);
	}*/
        //closesocket(sock);
}
